"""
realtime_monitor.py - 实时音频分类监测系统

功能：
1. 实时音频流捕获与处理
2. 低延迟音频分类预测
3. 线程安全的音频数据队列
4. 优雅的启动/停止控制
5. 详细的日志记录和错误处理

优化点：
- 添加类型注解
- 增强线程安全性
- 改进音频预处理
- 添加性能监控
- 更完整的资源清理

使用方法  python a08_realtime_monitor.py models/trained/model.pth --num_classes 2 --device cpu --sample_rate 48000 --chunk_size 4096

"""
import sys
import threading

import sounddevice as sd
import numpy as np
import torch
import time
import logging
from queue import Queue
from threading import Event, Thread
from typing import Optional, Tuple, Any
import argparse
import librosa  # 仅在需要时导入

from utils.audio_cnn import AudioCNN

# 重新配置（强制生效）
logging.basicConfig(
    force=True,  # 关键！覆盖现有配置
    level=logging.DEBUG,
    format='%(asctime)s [%(threadName)s] %(levelname)s: %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout)  # 强制输出到stderr
    ]
)

# 获取logger（必须在配置后调用）
logger = logging.getLogger(__name__)
logger.info("=== 日志系统初始化成功 ===")


class RealtimeAudioMonitor:
    """实时音频监测分类系统

    参数:
        model_path (str): 训练好的模型权重路径
        sample_rate (int): 音频采样率 (默认44100)
        chunk_size (int): 每次处理的音频帧数 (默认4096)
        num_classes (int): 分类类别数 (默认5)
        device (str): 计算设备 ('cuda'或'cpu')
    """

    def __init__(
            self,
            model_path: str,
            sample_rate: int = 44100,
            chunk_size: int = 4096,
            num_classes: int = 5,
            device: Optional[str] = None
    ):
        # 音频配置
        self.sample_rate = sample_rate
        self.chunk_size = chunk_size
        self.channels = 1  # 单声道输入

        # 线程控制
        self.audio_queue: Queue = Queue(maxsize=100)  # 防止内存爆炸
        self.processing_semaphore = threading.Semaphore(1)  # 控制处理速率
        self.stop_event = Event()
        self.processing_thread: Optional[Thread] = None

        # 设置计算设备
        self.device = torch.device(
            device if device else "cuda" if torch.cuda.is_available() else "cpu"
        )
        logger.info(f"使用计算设备: {self.device}")

        # 加载模型
        self.model = self._load_model(model_path, num_classes)

        # 性能监控
        self.processing_times = []
        self.frame_count = 0

        logger.info(f"实时监测器初始化完成 - 采样率: {sample_rate}Hz | 块大小: {chunk_size}")

        # 检查音频参数有效性
        assert self.chunk_size > 0, "块大小必须为正数"
        assert self.sample_rate in {8000, 16000, 44100, 48000}, "不支持的采样率"

        # 验证Librosa能否工作
        test_audio = np.random.rand(1024)
        try:
            librosa.feature.melspectrogram(y=test_audio, sr=self.sample_rate)
            print("Librosa测试通过")
        except Exception as e:
            print(f"Librosa测试失败: {e}")

    def _load_model(self, model_path: str, num_classes: int) -> torch.nn.Module:
        """加载训练好的PyTorch模型"""
        try:
            model = AudioCNN(num_classes=num_classes)

            # 加载模型权重 (支持多GPU训练保存的模型)
            state_dict = torch.load(model_path, map_location=self.device)
            if all(k.startswith('module.') for k in state_dict.keys()):
                state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}

            model.load_state_dict(state_dict)
            model.to(self.device)
            model.eval()

            logger.info(f"成功加载模型: {model_path}")
            return model
        except Exception as e:
            logger.error(f"加载模型失败: {str(e)}")
            raise

    def _audio_callback(
            self,
            indata: np.ndarray,
            frames: int,
            time: Any,
            status: sd.CallbackFlags
    ) -> None:
        """SoundDevice音频回调函数"""
        if status:
            logger.warning(f"音频流状态: {status}")

        # 非阻塞式放入队列
        if not self.audio_queue.full():
            self.audio_queue.put(indata.copy())
        else:
            logger.warning("音频队列已满，丢弃数据")

    def _preprocess(self, audio: np.ndarray) -> torch.Tensor:
        """预处理音频数据为模型输入格式

        参数:
            audio: 原始音频数据 (形状: [chunk_size])

        返回:
            形状为 [1, 1, n_mels, time] 的张量
        """
        try:
            # 确保音频长度足够（关键修复！）
            if len(audio) < 2048:  # 至少2048采样点
                audio = np.pad(audio, (0, 2048 - len(audio)))

            # 转换为梅尔频谱
            mel = librosa.feature.melspectrogram(
                y=audio,
                sr=self.sample_rate,
                n_mels=128,
                hop_length=512,
                n_fft=2048
            )

            # 统一尺寸为256x256
            if mel.shape[1] < 256:
                mel = np.pad(mel, ((0, 0), (0, 256 - mel.shape[1])))
            else:
                mel = mel[:, :256]

            # 转换为dB尺度
            mel_db = librosa.power_to_db(mel, ref=np.max)

            # 转换为PyTorch张量并添加批次和通道维度
            tensor = torch.FloatTensor(mel_db).unsqueeze(0).unsqueeze(0)

            return tensor.to(self.device)

        except Exception as e:
            logger.error(f"音频预处理失败: {str(e)}")
            return torch.zeros(1, 1, 128, 128)  # 返回安全值

    def _process_audio(self) -> None:
        """处理音频数据的线程函数"""
        logger.info("音频处理线程启动")

        while not self.stop_event.is_set():
            with self.processing_semaphore:  # 限制处理并发
                try:
                    if not self.audio_queue.empty():
                        start_time = time.time()

                        # 获取音频数据
                        audio = self.audio_queue.get().squeeze()

                        # 预处理和预测
                        inputs = self._preprocess(audio)
                        with torch.no_grad():
                            outputs = self.model(inputs)

                        # 获取预测结果
                        probs = torch.softmax(outputs, dim=1)
                        pred_class = torch.argmax(probs).item()
                        confidence = probs[0, pred_class].item()

                        # 记录性能
                        process_time = time.time() - start_time
                        self.processing_times.append(process_time)
                        self.frame_count += 1

                        # 打印结果 (可替换为其他输出方式)
                        logger.info(
                            f"预测类别: {pred_class} | "
                            f"置信度: {confidence:.2%} | "
                            f"延迟: {process_time:.3f}s"
                        )

                except Exception as e:
                    logger.error(f"音频处理异常: {str(e)}")
                    continue

        logger.info("音频处理线程停止")

    def start(self) -> None:
        """启动实时监测系统"""
        try:
            # # 打印所有可用设备
            # print("\n=== 可用音频设备 ===")
            # print(sd.query_devices())

            # 检查默认设备
            default_input = sd.default.device[0]
            logger.info(f"\n使用的输入设备ID: {default_input}")

            # 验证设备配置
            logger.info("\n尝试测试音频输入...")
            with sd.InputStream(
                    samplerate=self.sample_rate,
                    blocksize=self.chunk_size,
                    channels=self.channels
            ) as test_stream:
                logger.info(f"音频流信息: {test_stream}")
                logger.info("音频设备测试通过！")

            # 初始化音频流
            # self.stream = sd.InputStream(
            #     samplerate=self.sample_rate,
            #     blocksize=self.chunk_size,
            #     channels=self.channels,
            #     dtype='float32',
            #     callback=self._audio_callback
            # )
            self.stream = sd.InputStream(
                device=default_input,
                samplerate=self.sample_rate,
                blocksize=self.chunk_size,
                channels=self.channels,
                dtype='float32',
                callback=self._audio_callback
            )

            # 启动处理线程
            self.processing_thread = Thread(target=self._process_audio, daemon=True, name="AudioProcessor")
            self.stop_event.clear()
            self.processing_thread.start()

            # 启动音频流
            self.stream.start()
            logger.info("实时监测已启动 - 按Ctrl+C停止")

        except Exception as e:
            logger.error(f"启动失败: {str(e)}")
            self.stop()
            raise

    def stop(self) -> None:
        """停止实时监测系统"""
        logger.info("正在停止实时监测...")

        # 停止音频流
        if hasattr(self, 'stream'):
            self.stream.stop()
            self.stream.close()

        # 停止处理线程
        self.stop_event.set()
        if self.processing_thread and self.processing_thread.is_alive():
            self.processing_thread.join(timeout=1.0)

        # 打印性能统计
        if self.processing_times:
            avg_time = sum(self.processing_times) / len(self.processing_times)
            logger.info(
                f"处理统计 - 总帧数: {self.frame_count} | "
                f"平均延迟: {avg_time:.3f}s | "
                f"最大延迟: {max(self.processing_times):.3f}s"
            )


if __name__ == "__main__":
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='实时音频分类监测')
    parser.add_argument('model_path', type=str, help='模型权重路径')
    parser.add_argument('--sample_rate', type=int, default=44100,
                        help='音频采样率 (默认44100)')
    parser.add_argument('--chunk_size', type=int, default=4096,
                        help='音频块大小 (默认4096)')
    parser.add_argument('--num_classes', type=int, default=5,
                        help='分类类别数 (默认5)')
    parser.add_argument('--device', type=str,
                        help='计算设备 (cuda/cpu)，默认自动选择')


    args = parser.parse_args()

    logger.info("1 - 可用音频设备:")
    logger.info(sd.query_devices())

    try:
        # 让用户选择设备
        device_id = int(input("请输入要使用的输入设备ID: "))
        sd.default.device = device_id

        # 初始化监测器
        monitor = RealtimeAudioMonitor(
            model_path=args.model_path,
            sample_rate=args.sample_rate,
            chunk_size=args.chunk_size,
            num_classes=args.num_classes,
            device=args.device
        )


        # 启动监测
        monitor.start()

        # 主线程等待中断
        while True:
            time.sleep(0.1)

    except KeyboardInterrupt:
        logger.info("接收到中断信号")
    except Exception as e:
        logger.error(f"运行时错误: {str(e)}")
    finally:
        monitor.stop()
        logger.info("系统已安全关闭")